#include <vector>
#include <cmath>
#include <cassert>
#include "utool_AudioRT/Config.h"
#include "utool_AudioRT/Audio_RT_FFT_KF.h"
// 读取CSV
#include <fstream>
#include <sstream>
#include <vector>
#include <string>
#include <iostream>



// 创建Sys窗函数
std::vector<__fp16> asymmetric_Sys_windows(int hop_length, int win_length, int zero_d) {
    const int zero_s = win_length - 2 * hop_length;
    const int Bottom_C = win_length - hop_length - zero_d;

    // Generate hanning_window_analy_Half_2
    std::vector<__fp16> hanning_window_analy_Half_2;
    hanning_window_analy_Half_2.reserve(hop_length);
    for (int i = hop_length; i < 2 * hop_length; ++i) {
        const float x = static_cast<float>(i);
        const float val = std::sqrt(0.5f * (1.0f - std::cos(M_PI * x / hop_length)));
        hanning_window_analy_Half_2.push_back(static_cast<__fp16>(val));
    }
    // Generate A_N_Array and compute A_N
    const int A_N_start = win_length - 2 * hop_length - zero_d;
    const int A_N_end = win_length - hop_length - zero_d;
    std::vector<__fp16> A_N;
    A_N.reserve(hop_length);
    for (int i = A_N_start; i < A_N_end; ++i) {
        const float x = static_cast<float>(i);
        const float val = std::sqrt(0.5f * (1.0f - std::cos(M_PI * x / Bottom_C)));
        A_N.push_back(static_cast<__fp16>(val));
    }
    // Compute hanning_window_Sys_Ahead
    std::vector<__fp16> hanning_window_Sys_Ahead;
    hanning_window_Sys_Ahead.reserve(hop_length);
    for (int i = 0; i < hop_length; ++i) {
        const float sys_val = 0.5f * (1.0f - std::cos(M_PI * i / hop_length));
        const float a_n_val = static_cast<float>(A_N[i]);
        const float result = sys_val / a_n_val;
        hanning_window_Sys_Ahead.push_back(static_cast<__fp16>(result));
    }
    // Concatenate and pad
    std::vector<__fp16> hanning_window_Sys;
    hanning_window_Sys.reserve(hanning_window_Sys_Ahead.size() + hanning_window_analy_Half_2.size() + zero_s);
    hanning_window_Sys.insert(hanning_window_Sys.end(), hanning_window_Sys_Ahead.begin(), hanning_window_Sys_Ahead.end());
    hanning_window_Sys.insert(hanning_window_Sys.end(), hanning_window_analy_Half_2.begin(), hanning_window_analy_Half_2.end());
    hanning_window_Sys.insert(hanning_window_Sys.begin(), zero_s, static_cast<__fp16>(0));
    return hanning_window_Sys;
}

// 创建Analy窗函数
std::vector<__fp16> asymmetric_Analy_windows(int hop_length, int win_length, int zero_d) {
    const int Bottom_C = win_length - hop_length - zero_d;

    // Generate hanning_window_analy_Half
    const int Analy_Half_len = win_length - zero_d - hop_length;
    std::vector<__fp16> hanning_window_analy_Half;
    hanning_window_analy_Half.reserve(Analy_Half_len);
    for (int i = 0; i < Analy_Half_len; ++i) {
        const float x = static_cast<float>(i);
        const float val = std::sqrt(0.5f * (1.0f - std::cos(M_PI * x / Bottom_C)));
        hanning_window_analy_Half.push_back(static_cast<__fp16>(val));
    }

    // Generate hanning_window_analy_Half_2
    std::vector<__fp16> hanning_window_analy_Half_2;
    hanning_window_analy_Half_2.reserve(hop_length);
    for (int i = hop_length; i < 2 * hop_length; ++i) {
        const float x = static_cast<float>(i);
        const float val = std::sqrt(0.5f * (1.0f - std::cos(M_PI * x / hop_length)));
        hanning_window_analy_Half_2.push_back(static_cast<__fp16>(val));
    }

    // Concatenate and pad
    std::vector<__fp16> hanning_window_analy;
    hanning_window_analy.reserve(hanning_window_analy_Half.size() + hanning_window_analy_Half_2.size() + zero_d);
    hanning_window_analy.insert(hanning_window_analy.end(), hanning_window_analy_Half.begin(), hanning_window_analy_Half.end());
    hanning_window_analy.insert(hanning_window_analy.end(), hanning_window_analy_Half_2.begin(), hanning_window_analy_Half_2.end());
    hanning_window_analy.insert(hanning_window_analy.begin(), zero_d, static_cast<__fp16>(0));

    return hanning_window_analy;
}

// 应用窗口
std::vector<std::vector<__fp16>> apply_window_multiply(
    const std::vector<std::vector<__fp16>>& Time_Cache_Matrix,
    const std::vector<__fp16>& Anly_Windows)
{
    // 输入有效性检查
    assert(!Time_Cache_Matrix.empty());
    assert(!Anly_Windows.empty());
    assert(Time_Cache_Matrix[0].size() == Anly_Windows.size());

    // 预分配结果内存
    std::vector<std::vector<__fp16>> result;
    result.reserve(Time_Cache_Matrix.size());
    
    // 并行友好的逐元素乘法
    for (const auto& row : Time_Cache_Matrix) {
        std::vector<__fp16> multiplied_row;
        multiplied_row.reserve(row.size());
        
        for (size_t i = 0; i < row.size(); ++i) {
            // 使用 float 中间值避免精度损失
            const float val = static_cast<float>(row[i]) * static_cast<float>(Anly_Windows[i]);
            multiplied_row.push_back(static_cast<__fp16>(val));
        }
        
        result.push_back(std::move(multiplied_row));
    }
    
    return result;
}

// 原位处理最后一行数据
void multiply_last_n_rows_inplace(
    std::vector<std::vector<__fp16>>& matrix,
    const std::vector<__fp16>& window,
    size_t n_rows)
{
    const size_t start = matrix.size() > n_rows ? matrix.size() - n_rows : 0;
    for (size_t i = start; i < matrix.size(); ++i) {
        for (size_t j = 0; j < window.size(); ++j) {
            matrix[i][j] = static_cast<__fp16>(
                static_cast<float>(matrix[i][j]) * 
                static_cast<float>(window[j]));
        }
    }
}



// 应用分析窗口结果，只处理末尾的两帧，并且只特定处理最后一帧倒数第二个Hop块，以及倒数第二帧最后一个Hop块，把两个块相加
std::vector<__fp16> Generate_SysResult(
    const std::vector<std::vector<__fp16>>& matrix,
    const std::vector<__fp16>& window,
    size_t n_rows)
{
    std::vector<std::vector<__fp16>> result;
    const size_t start_row = matrix.size() > n_rows ? matrix.size() - n_rows : 0;

    std::vector<__fp16> row_result;
    // row_result.resize(HOP_SIZE);
    for (size_t j = 0; j < HOP_SIZE; ++j) {
        const float val_end2 = static_cast<float>(matrix[start_row][window.size() - HOP_SIZE - 1 + j]) 
                        * static_cast<float>(window[window.size() - HOP_SIZE - 1 + j]);
        const float val_end = static_cast<float>(matrix[start_row+1][window.size() - 2*HOP_SIZE - 1 + j])
        * static_cast<float>(window[window.size() - 2*HOP_SIZE - 1 + j]);
        row_result.push_back(static_cast<__fp16>(val_end2 + val_end));
    }

    
    return row_result;
}

// 输入帧和幅度谱，计算这一帧的能量(直接算总和，因为已经是能量了)
float RT_Energe_Frame(std::vector<std::vector<__fp16>>& mag,
    size_t Frame)
{
    float Energe_Value = 0;
    for (size_t k = 0; k < FFT_BINS; ++k) {
        Energe_Value = Energe_Value + mag[Frame][k];
    }
    return Energe_Value;
}



// 输入幅度谱，只处理这个谱的最后两行，计算SPL，然后计算插值增益，最后得到增益后谱
void RT_WDRC(std::vector<std::vector<__fp16>>& mag,
        std::vector<std::vector<__fp16>>& WDRC_mag,
        std::vector<std::vector<__fp16>>& WDRC_Table
    ) 
{
    WDRC_mag.resize(2, std::vector<__fp16>(FFT_BINS)); // 固定就是最后2帧的IFFT
    size_t num_frames = mag.size(); // 动态获取mag的行数
    size_t Gain_Index = 0; // 设置增益表格的序号
    float inputSPL_Frame = 0; // 整帧声压级
    float inputSPL = 0;
    float Energe_Value = 0;
    for (size_t i = num_frames - 2; i < num_frames; ++i) {
        // 进行WDRC
        // 先判断这一帧的总能量，如果这一帧的总能量小于某一程度，则直接静音
        Energe_Value = RT_Energe_Frame(mag,i);
        // 计算当前帧的声压级
        inputSPL_Frame = 20 * std::log10(Energe_Value) + 120;
        if(inputSPL_Frame <= 93){
            for (size_t k = 0; k < FFT_BINS; ++k) {
                WDRC_mag[i - num_frames + 2][k] = static_cast<__fp16>(0);
            }
        }
        else{
            for (size_t k = 0; k < FFT_BINS; ++k) {
                //计算输入SPL
                float max_mag = 0;
                //计算幅值
                // max_mag = std::max(0.000000001f, static_cast<float>(mag[i][k]));
                if (mag[i][k] > 0.0000001f)
                {
                    //根据各通道校准值转SPL
                    inputSPL = 20 * std::log10(mag[i][k]) + 120; // Spl_Max_dB[k]
                    // 根据计算的各通道SPL数值计算增益数值
                    // 除以2并进行四舍五入取整
                    float temp = inputSPL / 2.0f;
                    int Gain_Index = static_cast<int>(std::round(temp)); // 设置增益表格的序号
                    // 确保值在0到55之间
                    if (Gain_Index < 0) {
                        Gain_Index = 0;
                    } else if (Gain_Index > 55) {
                        Gain_Index = 55;
                    }
                    // float gainValue = WDRC_Table[k][Gain_Index]; // 读取WDRC表格进而读取增益值
                    float gainValue = 20; // 读取WDRC表格进而读取增益值
                    // 利用对数恒等式：10^(x/20) = exp(x * ln(10)/20)
                    constexpr float ln10_div_20 = 0.115129f;  // ln(10)/20 ≈ 0.115129
                    float tmp = std::exp((inputSPL + gainValue - Spl_Max_dB[k]) * ln10_div_20);
                    WDRC_mag[i - num_frames + 2][k] = static_cast<__fp16>(tmp);
                }
                else
                {
                    WDRC_mag[i - num_frames + 2][k] = static_cast<__fp16>(0);
                }
            }
        }
        
    }
}

void RT_FC(std::vector<std::vector<__fp16>>& mag,
        std::vector<std::vector<__fp16>>& FC_mag,
        std::vector<std::vector<__fp16>>& FC_Table)
{
    FC_mag.resize(2, std::vector<__fp16>(FFT_BINS)); // 固定就是最后2帧的IFFT
    size_t num_frames = mag.size(); // 动态获取mag的行数
    size_t inner = FC_Table.size();          // FFT_BINS
    size_t cols = FC_Table[0].size();        // FFT_BINS

    for (size_t i = num_frames - 2; i < num_frames; ++i) {
        for (size_t j = 0; j < cols; ++j) {
            __fp16 sum = 0;
            for (size_t k = 0; k < inner; ++k) {
                sum += mag[i][k] * FC_Table[k][j];
            }
            FC_mag[i - num_frames + 2][j] = sum;
        }
    }
}

void RT_FC_CSR(
    const std::vector<std::vector<__fp16>>& WDRC_mag,  // dense matrix 2 x FFT_BINS
    const std::vector<__fp16>& values,                 // FC_Table 非零元素
    const std::vector<int>& col_idx,                    // 对应列索引
    const std::vector<int>& row_ptr,                    // 每行起始位置
    std::vector<std::vector<__fp16>>& FC_mag) {

    size_t num_frames = WDRC_mag.size(); // 动态获取mag的行数
    FC_mag.assign(2, std::vector<__fp16>(FFT_BINS, 0)); // 固定就是最后2帧的IFFT

    for (size_t i = num_frames - 2; i < num_frames; ++i) {          // WDRC_mag 行
        for (int row = 0; row < FFT_BINS; ++row) {  // FC_Table 行
            __fp16 w = WDRC_mag[i][row];
            for (int idx = row_ptr[row]; idx < row_ptr[row + 1]; ++idx) {
                int col = col_idx[idx];
                __fp16 val = values[idx];
                FC_mag[i - num_frames + 2][col] += w * val;
            }
        }
    }
}



// 用于转换为CSR格式
void convertToCSR(std::vector<std::vector<__fp16>>& FC_Table,
                  std::vector<__fp16>& values,
                  std::vector<int>& col_idx,
                  std::vector<int>& row_ptr,
                  __fp16 threshold = 1e-5)  // 忽略小于该值的元素
{
    int rows = FC_Table.size();
    if (rows == 0) return;
    int cols = FC_Table[0].size();

    values.clear();
    col_idx.clear();
    row_ptr.clear();
    row_ptr.push_back(0); // 第一行从位置0开始

    for (int i = 0; i < rows; ++i) {
        for (int j = 0; j < cols; ++j) {
            __fp16 val = FC_Table[i][j];
            if (std::abs(static_cast<float>(val)) > threshold) {
                values.push_back(val);
                col_idx.push_back(j);
            }
        }
        row_ptr.push_back(static_cast<int>(values.size()));  // 下一行起始
    }
}


void Read_WDRC_Table(std::vector<std::vector<__fp16>>& WDRC_Table)
{
    const std::string filename = "WDRC_Table.csv";
    // 打开CSV文件
    std::ifstream file(filename);
    if (!file.is_open()) {
        std::cerr << "错误: 无法打开文件 " << filename << std::endl;
    }
    std::string line;
    int lineCount = 0;
    
    // 逐行读取文件
    while (std::getline(file, line)) {
        lineCount++;
        // 跳过第一行（表头）
        if (lineCount == 1) continue;
        
        std::vector<__fp16> row;
        std::istringstream lineStream(line);
        std::string cell;
        
        // 解析每行的逗号分隔值
        while (std::getline(lineStream, cell, ',')) {
            try {
                // 将字符串转换为float后再转换为__fp16
                __fp16 value = static_cast<__fp16>(std::stof(cell));
                row.push_back(value);
            } catch (const std::exception& e) {
                std::cerr << "警告: 第" << lineCount << "行转换错误: " << e.what() << std::endl;
            }
        }
        // 将行数据添加到表格
        if (!row.empty()) {
            WDRC_Table.push_back(row);
        }
    }

    file.close();
    
    // 显示成功信息
    if (WDRC_Table.empty()) {
        std::cout << "警告: 表格数据为空!" << std::endl;
    } else {
        std::cout << "成功读取 " << WDRC_Table.size() << " 行数据" << std::endl;
        std::cout << "每行 " << WDRC_Table[0].size() << " 列" << std::endl;
        std::cout << "示例数据 (第2行首2个值): " << std::endl;
        for (size_t i = 0; i < 2 && i < WDRC_Table[0].size(); i++) {
            std::cout << static_cast<float>(WDRC_Table[0][i]) << " ";
        }
        std::cout << std::endl;
    }

}

void Read_FC_Table(std::vector<std::vector<__fp16>>& FC_Table)
{
    const std::string filename = "FC_Table.csv";
    // 打开CSV文件
    std::ifstream file(filename);
    if (!file.is_open()) {
        std::cerr << "错误: 无法打开文件 " << filename << std::endl;
    }
    std::string line;
    int lineCount = 0;
    
    // 逐行读取文件
    while (std::getline(file, line)) {
        lineCount++;
        // 跳过第一行（表头）

        std::vector<__fp16> row;
        std::istringstream lineStream(line);
        std::string cell;
        
        // 解析每行的逗号分隔值
        while (std::getline(lineStream, cell, ',')) {
            try {
                // 将字符串转换为float后再转换为__fp16
                __fp16 value = static_cast<__fp16>(std::stof(cell));
                row.push_back(value);
            } catch (const std::exception& e) {
                std::cerr << "警告: 第" << lineCount << "行转换错误: " << e.what() << std::endl;
            }
        }
        // 将行数据添加到表格
        if (!row.empty()) {
            FC_Table.push_back(row);
        }
    }

    file.close();
    
    // 显示成功信息
    if (FC_Table.empty()) {
        std::cout << "警告: 表格数据为空!" << std::endl;
    } else {
        std::cout << "成功读取 " << FC_Table.size() << " 行数据" << std::endl;
        std::cout << "每行 " << FC_Table[0].size() << " 列" << std::endl;
        std::cout << "示例数据 (第2行首2个值): " << std::endl;
        for (size_t i = 0; i < 2 && i < FC_Table[0].size(); i++) {
            std::cout << static_cast<float>(FC_Table[0][i]) << " ";
        }
        std::cout << std::endl;
    }

}

// 计算RMS值的函数
float calculateRMS(const std::vector<float>& audioData) {
    if (audioData.empty()) {
        return 0.0f;
    }
    
    double sumSquares = 0.0;
    
    // 计算平方和
    for (float sample : audioData) {
        sumSquares += static_cast<double>(sample) * static_cast<double>(sample);
    }
    
    // 计算均值并开方
    double meanSquares = sumSquares / audioData.size();
    return static_cast<float>(std::sqrt(meanSquares));
}